import numpy as np
import tensorflow as tf
from keras import layers
from keras.layers import Input, Add, Dense, Activation, ZeroPadding2D, BatchNormalization, Flatten, Conv2D, AveragePooling2D, MaxPooling2D, GlobalMaxPooling2D
from keras.models import Model, load_model
from keras.preprocessing import image
from keras.utils import layer_utils
from keras.utils.data_utils import get_file
from keras.applications.imagenet_utils import preprocess_input
# import pydot
# from IPython.display import SVG
from keras.utils.vis_utils import model_to_dot
from keras.utils import plot_model
from keras.initializers import glorot_uniform
import scipy.misc
from matplotlib.pyplot import imshow


import keras.backend as K
from math import sqrt
from numpy import concatenate

def convolutional_block(input_tensor, filters):
    filter1, filter2, filter3 = filters

    x = Conv2D(filter1,(1,1),strides=1)(input_tensor)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('tanh')(x)

    x = Conv2D(filter2,(3,1),strides=1,padding='same')(x)
    x = Conv2D(filter2,(1,3),strides=1,padding='same')(x)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('selu')(x)

    x = Conv2D(filter3,(1,1),strides=1)(x)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('selu')(x)

    y = Conv2D(filter3,(1,1),strides=1)(input_tensor)
    y = BatchNormalization(axis=-1)(y)
    y = Activation('elu')(y)

    # out = merge([x,y],mode='sum')
    #out=concatenate([x,y],axis=3)
    out=Add()([x,y])
    z = Activation('tanh')(out)

    return z

def identity_block(input_tensor, filters):


    filter1, filter2, filter3 = filters

    x = Conv2D(filter1,(1,1),strides=1)(input_tensor)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('tanh')(x)

    x = Conv2D(filter2,(3,1),strides=1,padding='same')(x)
    x = Conv2D(filter2,(1,3),strides=1,padding='same')(x)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('selu')(x)

    x = Conv2D(filter3,(1,1),strides=1)(x)
    x = BatchNormalization(axis=-1)(x)
    x = Activation('selu')(x)

    # y = Conv2D(filter3,(1,1),strides=1)(input_tensor)
    # y = BatchNormalization(axis=-1)(y)
    # y = Activation('elu')(y)

    # out = merge([x,input_tensor],mode='sum')
    #out=concatenate([x,y],axis=3)
    out=Add()([x,input_tensor])
    z = Activation('tanh')(out)
    return z

def ResNet50(input_shape = (64, 64, 3), classes = 6):
 
    # Define the input as a tensor with shape input_shape
    X_input = Input(input_shape)

    
    # Zero-Padding
    X = ZeroPadding2D((3, 3))(X_input)
    
    # Stage 1
    X = Conv2D(filters=64, kernel_size=(7, 7), strides=(2, 2), name="conv",
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name="bn_conv1")(X)
    X = Activation("relu")(X)
    X = MaxPooling2D(pool_size=(3, 3), strides=(2, 2))(X)

    # Stage 2
    X = convolutional_block(X, filters=[64, 64, 256])
    X = identity_block(X, filters=[64, 64, 256])
    X = identity_block(X, filters=[64, 64, 256])
    ### START CODE HERE ###

    # Stage 3 (≈4 lines)
    # The convolutional block uses three set of filters of size [128,128,512], "f" is 3, "s" is 2 and the block is "a".
    # The 3 identity blocks use three set of filters of size [128,128,512], "f" is 3 and the blocks are "b", "c" and "d".
    X = convolutional_block(X, filters=[128, 128, 512])
    X = identity_block(X,   filters=[128, 128, 512]  )
    X = identity_block(X,   filters=[128, 128, 512]  )
    X = identity_block(X,   filters=[128, 128, 512]  )
    
    # Stage 4 (≈6 lines)
    # The convolutional block uses three set of filters of size [256, 256, 1024], "f" is 3, "s" is 2 and the block is "a".
    # The 5 identity blocks use three set of filters of size [256, 256, 1024], "f" is 3 and the blocks are "b", "c", "d", "e" and "f".
    X = convolutional_block(X,   filters=[256, 256, 1024])
    X = identity_block(X,   filters=[256, 256, 1024]  )
    X = identity_block(X,   filters=[256, 256, 1024]  )
    X = identity_block(X,   filters=[256, 256, 1024]  )
    X = identity_block(X,   filters=[256, 256, 1024]  )
    X = identity_block(X,   filters=[256, 256, 1024]  )
    

    # Stage 5 (≈3 lines)
    # The convolutional block uses three set of filters of size [512, 512, 2048], "f" is 3, "s" is 2 and the block is "a".
    # The 2 identity blocks use three set of filters of size [256, 256, 2048], "f" is 3 and the blocks are "b" and "c".
    X = convolutional_block(X,   filters=[512, 512, 2048])
    X = identity_block(X,   filters=[512, 512, 2048])
    X = identity_block(X,   filters=[512, 512, 2048])
    
    # filters should be [256, 256, 2048], but it fail to be graded. Use [512, 512, 2048] to pass the grading
    

    # AVGPOOL (≈1 line). Use "X = AveragePooling2D(...)(X)"
    # The 2D Average Pooling uses a window of shape (2,2) and its name is "avg_pool".
    X = AveragePooling2D(pool_size=(2, 2), padding="same")(X)
    
    ### END CODE HERE ###

    # output layer
    X = Flatten()(X)
    X = Dense(classes, activation="softmax", name="fc"+str(classes), kernel_initializer=glorot_uniform(seed=0))(X)
    
    
    # Create model
    model = Model(inputs=X_input, outputs=X, name="ResNet50")

    return model
